<template>
  <div>
    <n-form :model="sdImg2ImgPrompt" :rules="rules" ref="formRef">

      <n-form-item :show-label="false" label="参考图片" path="negativePrompt">
        <n-card title="参考图片" bordered>
          <template #header-extra>
            <n-checkbox v-model:checked="drawingEnable" @update:checked="enablePartPaint">
              启用局部绘制
            </n-checkbox>
            <n-button v-if="drawingEnable" size="small" type="primary" style="margin-left: 5px;" @click="canvasUndo">
              撤销
            </n-button>
            <n-button v-if="drawingEnable" size="small" type="primary" style="margin-left: 5px;" @click="canvasRedo">
              回撤
            </n-button>
            <n-button size="small" type="warning" style="margin-left: 5px;" @click="resetImg"> 清除</n-button>

          </template>
          <n-space justify="center">
            <div class="canvas-wrapper">
              <canvas
                  ref="canvasRef"
                  :width="canvasWidth"
                  :height="canvasHeight"
                  @mousedown="startDrawing" @mousemove="draw" @mouseup="stopDrawing"></canvas>
              <img
                  ref="img2img_imgRef"
                  :style="{
                      width: canvasWidth + 'px',
                      height: canvasHeight + 'px'
                      }"
                  v-if="sdImg2ImgPrompt.image"
                  :src="sdImg2ImgPrompt.image"
                  alt=""
                  @load="handleImgLoad"/>
            </div>
            <n-empty description="未上传图片" v-if="!sdImg2ImgPrompt.image">
              <template #extra>
                <n-upload
                    :custom-request="customRequest"
                >
                  <n-button>上传图片</n-button>
                </n-upload>
              </template>
            </n-empty>
          </n-space>
        </n-card>
      </n-form-item>
      <n-button block type="primary" style="margin-bottom: 5px" :loading="interrogateLoading" @click="interrogate">图片标签反推</n-button>
      <n-form-item label="正向提示词" path="prompt">
        <n-input v-model:value="sdImg2ImgPrompt.prompt" style="height: 100px" type="textarea"
                 placeholder="请输出提示词"/>


      </n-form-item>


      <n-form-item label="反向提示词" path="negativePrompt">
        <n-input placeholder="请输入反向提示词" type="textarea" v-model:value="sdImg2ImgPrompt.negativePrompt"/>
      </n-form-item>
      <n-form-item label="尺寸" path="width">
        <n-input-number v-model:value="sdImg2ImgPrompt.height" style="width: 50%; ">
          <template #prefix>
            高度
          </template>
        </n-input-number>
        <n-input-number v-model:value="sdImg2ImgPrompt.width" style="width: 50%;margin-left: 10px;">
          <template #prefix>
            长度
          </template>
        </n-input-number>
      </n-form-item>
      <n-row>
        <n-col :span="8">
          <n-form-item label="随机种子" path="seed">
            <n-input-number v-model:value="sdImg2ImgPrompt.seed"/>
            <n-button type="primary" @click="sdImg2ImgPrompt.seed = Math.round(Math.random() * 1000000)">随机</n-button>
          </n-form-item>
        </n-col>
        <n-col :span="8">
          <n-form-item label="采样器" path="sampler">
            <n-select
                :options="samperOptions"
                v-model:value="sdImg2ImgPrompt.samplerName"
                placeholder="请选择"
                style="width: 100%;margin-left: 10px;"
            />
          </n-form-item>
        </n-col>
        <n-col :span="8">
          <n-form-item label="采样迭代步数" path="defaultSteps">
            <n-input-number v-model:value="sdImg2ImgPrompt.steps" style="width: 100%;margin-left: 10px;"/>
          </n-form-item>
        </n-col>
      </n-row>

      <n-row>
        <n-col :span="8">
          <n-form-item label="提示词相关性" path="CfgScale">
            <n-input-number v-model:value="sdImg2ImgPrompt.cfgScale" style="width: 100%"/>
          </n-form-item>
        </n-col>
        <n-col :span="8">
          <n-form-item label="重绘幅度" path="denoisingStrength">
            <n-input-number v-model:value="sdImg2ImgPrompt.denoisingStrength" style="width: 100%;margin-left: 10px;"/>
          </n-form-item>
        </n-col>
        <n-col :span="8">
          <n-form-item label="生成数量" path="Batch Size">
            <n-input-number v-model:value="sdImg2ImgPrompt.batchSize" style="width: 100%;margin-left: 10px;"/>
          </n-form-item>
        </n-col>
      </n-row>
    </n-form>
    <n-drawer v-model:show="promptTemplateShow" :width="promptTemplateWidth">
      <n-drawer-content title="提示词选择器" closable>
        <n-tabs animated type="segment">
          <n-tab-pane name="tag" tab="标签">
            <promptTemplate :prompt="sdTxt2ImgPrompt.prompt" :negative-prompt="sdTxt2ImgPrompt.negativePrompt"
                            @loadPrompt-event="loadPromptEvent"></promptTemplate>
          </n-tab-pane>
          <n-tab-pane name="pic" tab="照片">
            <pic-template @loadConfig-event="loadConfigEvent"></pic-template>
          </n-tab-pane>
          <n-tab-pane name="present" tab="预设">
            <pic-template @loadConfig-event="loadConfigEvent"></pic-template>
          </n-tab-pane>
        </n-tabs>
      </n-drawer-content>
    </n-drawer>
  </div>
</template>

<script setup lang="ts">
import {onMounted, reactive, ref, watch} from 'vue'
import {UploadCustomRequestOptions, useMessage} from 'naive-ui'
import {sdControllerApi} from '@/api/sd/sdController.api'
import promptTemplate from "@/components/prompt/promptTemplate.vue"
import PicTemplate from "@/components/prompt/picTemplate.vue"
import {fileController} from '@/api/file/fileController'
import {storeToRefs} from "pinia";
import {useUserStore} from "@/store/modules/useUserStore";
import {store} from "@/store";
import {storage} from "@/hooks/useStorage";


const rules = {
  prompt: {
    required: true,
    message: '请输入提示词',
    trigger: 'blur',
  },
  mobile: {
    required: true,
    message: '请输入联系电话',
    trigger: 'input',
  },
};
const promptTemplateShow = ref(false)
const formRef: any = ref(null);
const message = useMessage();
const imageSrc = ref(); // 默认图片的URL
const taskStatus = ref("waiting");
const promptTemplateWidth = ref(0)
const canvasRef = ref(null)
const ctx = ref(null)
const drawing = ref(false)
const drawingEnable = ref(false)
const canvasHistory = ref([])
const drawStep = ref(-1)
const img2img_imgRef = ref(null)
const {authorization,tokenKey} = storeToRefs(useUserStore(store));
const token = authorization?.value ?? storage.get(tokenKey?.value as string);


 const interrogateLoading = ref(false)
const selectInterrogator = ref('wd14-vit-v2-git')

onMounted(() => {
  // 设置当前drawerWidth 宽度
  promptTemplateWidth.value = window.innerWidth - 300;
  window.addEventListener('resize', handleResize);
})

const props = defineProps({
  resourceId: {
    type: String,
    default: '',
    required: true
  },
  defaultSdImg2ImgPrompt: {
    type: Object,
    default: () => {
      return {}
    }
  },
  samperOptions : {
    type: Array,
    default: () => {
      return []
    }
  }
})


watch(props.defaultSdImg2ImgPrompt, (newVal) => {
  sdImg2ImgPrompt.prompt = newVal.prompt
  sdImg2ImgPrompt.negativePrompt = newVal.defaultNegativePrompt
  sdImg2ImgPrompt.seed = newVal.defaultSeed
  sdImg2ImgPrompt.samplerName = newVal.defaultSamplerName
  sdImg2ImgPrompt.batchSize = newVal.defaultBatchSize
  sdImg2ImgPrompt.steps = newVal.defaultSteps
  sdImg2ImgPrompt.width = newVal.defaultWidth
  sdImg2ImgPrompt.height = newVal.defaultHeight
  sdImg2ImgPrompt.promptEn = newVal.promptEn
  sdImg2ImgPrompt.cfgScale = newVal.defaultCfgScale
  sdImg2ImgPrompt.savePath = newVal.savePath
  sdImg2ImgPrompt.restoreFaces = newVal.defaultRestoreFaces
  sdImg2ImgPrompt.denoisingStrength = newVal.defaultDenoisingStrength
  console.log(newVal.img2imgFileId)
  if (newVal.img2imgFileId) {
    loadImg(newVal.img2imgFileId)
  }
}, {deep: true})

// 窗口宽度变化
const handleResize = () => {
  promptTemplateWidth.value = window.innerWidth - 300;
}


const sdImg2ImgPrompt: SD.Img2imgInput = reactive({
  prompt: '',
  negativePrompt: '',
  seed: -1,
  samplerName: 'DPM++ 2M',
  batchSize: 1,
  steps: 20,
  width: 512,
  height: 512,
  promptEn: '',
  cfgScale: 7,
  savePath: '',
  image: '',
  denoisingStrength: 0.75
});

const samplerOptions = [
  {
    label: 'DPM++ 2M Karras',
    value: 'DPM++ 2M Karras'
  },
  {
    label: 'k_heun',
    value: 'k_heun'
  },
  {
    label: "k_euler",
    value: 'k_euler'
  },
  {
    label: "k_euler_a",
    value: 'k_euler_a'
  },
  {
    label: "k_dpm_2",
    value: 'k_dpm_2'
  },
  {
    label: "k_dpm_2_a",
    value: 'k_dpm_2_a'
  },
  {
    label: "k_dpm_fast",
    value: 'k_dpm_fast'
  },
  {
    label: "k_dpm_adaptive",
    value: 'k_dpm_adaptive'
  },
  {
    label: "k_dpmpp_2s_a",
    value: 'k_dpmpp_2s_a'
  },
  {
    label: "k_dpmpp_2m",
    value: 'k_dpmpp_2m'
  },
  {
    label: "dpmsolver",
    value: 'dpmsolver'
  },
  {
    label: "k_dpmpp_sde",
    value: 'k_dpmpp_sde'
  },
  {
    label: "DDIM",
    value: 'DDIM'
  }
];


const formSubmit = async () => {
  const res = await sdControllerApi.generateImg2Img(sdImg2ImgPrompt)
  if (res.code === 200) {
    const data = res.data;
    if (data.code === 1) {
      checkTask(data.result);
    }
  } else {
    message.error(res.message);
  }

}

let timeoutId = ref();// 保存当前的 setTimeout 的 id
const checkTask = async (taskId: string) => {
  console.log("checkTask", taskId)
  const response = await sdControllerApi.checkTask(taskId);
  if (response.code === 200) {
    const data = response.data;
    taskStatus.value = data.status;
    if (data.status !== "SUCCESS") {
      timeoutId.value = setTimeout(() => {
        checkTask(taskId);
      }, 1000); // 每隔0.5秒调用一次 checkTask
    } else {
      clearTimeout(timeoutId.value); // 当任务完成时停止调用 checkTask
      imageSrc.value = "http://127.0.0.1:8066/filetransfer/preview?promptId=" + taskId;

    }
  } else {
    clearTimeout(timeoutId.value); // 当任务完成时停止调用 checkTask
    imageSrc.value = "http://127.0.0.1:8066/filetransfer/preview?promptId=" + taskId;
  }

}
// 设置canvas宽高
const canvasWidth = ref(600)
const canvasHeight = ref(600)

const loadImg = (imgFileId:string) =>{
  // 加载图片
  fileController.getFileBase64(token,imgFileId).then(res => {
    if (res.code === 200) {
      sdImg2ImgPrompt.image = `data:image/jpeg;base64,`+res.data
      ctx.value = canvasRef.value.getContext('2d')
    } else {
      message.error(res.message);
    }
  })
}

const handleImgLoad =  (e: any) => {
  const image = e.target;
  let image_new = new Image()
  image_new.src = image.src
  image_new.onload = function () {
    sdImg2ImgPrompt.width = image_new.width
    sdImg2ImgPrompt.height = image_new.height
    // 计算图片宽高比例
    let imgRatio = image_new.width / image_new.height
    if (image_new.width > 600
        || image_new.height > 600) {
      // 计算缩放后的宽高
      if (imgRatio > 1) {
        canvasWidth.value = 600
        canvasHeight.value = image_new.height / (image_new.width / 600)
      } else {
        // 图片比较宽
        canvasWidth.value = image_new.width / (image_new.height / 600)
        canvasHeight.value = 600
      }
    } else {
      canvasWidth.value = image_new.width
      canvasHeight.value = image_new.height
    }
  }
}


const customRequest = (uploadOption: UploadCustomRequestOptions) => {
  const reader = new FileReader();
  //@ts-ignore
  reader.readAsDataURL(uploadOption.file.file);
  console.log(uploadOption.file.file)
  reader.onloadend = () => {
    ctx.value = canvasRef.value.getContext('2d')
    let image = new Image()
    image.src = reader.result
    //@ts-ignore
    sdImg2ImgPrompt.image = reader.result
    console.log(reader.result)

  };
}

const loadPromptEvent = (prompt: string, negativePrompt: string) => {
  sdImg2ImgPrompt.prompt = prompt;
  sdImg2ImgPrompt.negativePrompt = negativePrompt;
  promptTemplateShow.value = false;
}




const sortFunction = (arr) => {
  return arr.sort((a, b) => {
    if (a.song > b.song) return 1;
    else return -1;
  });
}

const interrogate = () => {
  interrogateLoading.value = true
  const data = {
    resourceId: props.resourceId,
    image: sdImg2ImgPrompt.image,
    model: selectInterrogator.value,
    threshold: 0.35
  }
  sdControllerApi.interrogate(data).then(res => {
    if (res.code === 200) {
      sdImg2ImgPrompt.prompt = res.data
    } else {
      message.error(res.message);
    }
  }).finally(() => {
    interrogateLoading.value = false
  })
}


const startDrawing = (e: any) => {
  if (!ctx.value) return;
  ctx.value.strokeStyle = 'white'
  ctx.value.fillStyle = 'white'
  drawing.value = true
  if (drawing.value && drawingEnable.value) {
    ctx.value.beginPath()
    ctx.value.moveTo(e.clientX - canvasRef.value.getBoundingClientRect().left, e.clientY - canvasRef.value.getBoundingClientRect().top)
  }
}


const draw = (e: any) => {
  if (!drawing.value || !drawingEnable.value) return;
  drawCircle(e)
}

const drawCircle = (e: any) => {
  ctx.value.beginPath();
  ctx.value.fillStyle = '#2196f3';
  ctx.value.strokeStyle = '#2196f3';
  ctx.value.globalAlpha = 0.1
  ctx.value.arc(e.clientX - canvasRef.value.getBoundingClientRect().left, e.clientY - canvasRef.value.getBoundingClientRect().top, 10, 0, Math.PI * 2);
  ctx.value.stroke()
  ctx.value.fill()

  ctx.value.beginPath();
  ctx.value.fillStyle = '#2196f3';
  ctx.value.strokeStyle = '#2196f3';
  ctx.value.globalAlpha = 1
  ctx.value.arc(e.clientX - canvasRef.value.getBoundingClientRect().left, e.clientY - canvasRef.value.getBoundingClientRect().top, 10, 0, Math.PI * 2);
  ctx.value.stroke()
  ctx.value.fill()
  ctx.value.closePath()
}


const stopDrawing = () => {
  drawing.value = false
  ctx.value.closePath()
  drawStep.value = drawStep.value + 1
  if (drawStep.value < canvasHistory.value.length) {
    canvasHistory.value.length = drawStep.value; // 截断数组
  }
  canvasHistory.value.push(canvasRef.value.toDataURL())
}


const canvasUndo = () => {
  if (drawStep.value > 0) {
    drawStep.value = drawStep.value - 1
    let canvasPic = new Image();
    canvasPic.src = canvasHistory.value[drawStep.value]
    canvasPic.addEventListener('load', () => {
      ctx.value.clearRect(0, 0, canvasWidth.value, canvasHeight.value)
      console.log(canvasPic)
      ctx.value.drawImage(canvasPic, 0, 0)
    });
  } else {
    ctx.value.clearRect(0, 0, canvasWidth.value, canvasHeight.value)
  }
}

// 反撤销方法
const canvasRedo = () => {
  if (drawStep.value < canvasHistory.value.length - 1) {
    drawStep.value = drawStep.value + 1;
    let canvasPic = new Image();
    canvasPic.src = canvasHistory.value[drawStep.value]
    canvasPic.addEventListener('load', () => {
      ctx.value.clearRect(0, 0, canvasWidth.value, canvasHeight.value);
      ctx.value.drawImage(canvasPic, 0, 0);
    });
  }
}


const getBase64Mask = async () => {
  const imageData = ctx.value.getImageData(0, 0, canvasWidth.value, canvasHeight.value)
  const imageDataContent = imageData.data;
  // 修改图像数据中的像素颜色
  for (let i = 0; i < imageDataContent.length; i += 4) {
    // 判断透明度是否小于阈值
    if (imageDataContent[i + 3] < 128) {
      // 将透明部分设置为黑色
      imageDataContent[i] = 0;
      imageDataContent[i + 1] = 0;
      imageDataContent[i + 2] = 0;
      imageDataContent[i + 3] = 255;
    } else {
      // 将线条部分设置为白色
      imageDataContent[i] = 255;
      imageDataContent[i + 1] = 255;
      imageDataContent[i + 2] = 255;
      imageDataContent[i + 3] = 255;
    }
  }
  // 将修改后的图像数据导出为图片格式
  const exportCanvas = document.createElement('canvas');
  exportCanvas.width = canvasWidth.value
  exportCanvas.height = canvasHeight.value
  const exportCtx = exportCanvas.getContext('2d');
  exportCtx.putImageData(imageData, 0, 0);
  // mask放大到和原图一致
  return new Promise((resolve, reject) => {
    // scaleImage(exportCanvas.toDataURL(), sdImg2ImgPrompt.width, sdImg2ImgPrompt.height)
    const image = new Image()
    image.src = exportCanvas.toDataURL()
    image.onload = () => {
      const canvas = document.createElement('canvas')
      canvas.width = sdImg2ImgPrompt.width
      canvas.height = sdImg2ImgPrompt.height
      const ctx_canvas = canvas.getContext('2d')
      ctx_canvas.drawImage(image, 0, 0, sdImg2ImgPrompt.width, sdImg2ImgPrompt.height);
      resolve(canvas.toDataURL())
    }
  })

}

// const scaleImage = (src, width, height) => {
//     const image = new Image();
//     image.src = src;
//     image.onload = () => {
//       const canvas = document.createElement('canvas');
//       canvas.width = width
//       canvas.height = height
//       const ctx_canvas = canvas.getContext('2d');
//       ctx_canvas.drawImage(image, 0, 0, width, height);
//       sdImg2ImgPrompt.mask = canvas.toDataURL();
//       resolve(canvas.toDataURL())
//     }
// }


const resetImg = () => {
  sdImg2ImgPrompt.image = ''
  ctx.value.clearRect(0, 0, canvasWidth.value, canvasHeight.value)
  canvasHistory.value = []
}

const getImg2ImgPrompt = async () => {
  if (drawingEnable.value) {
    const response = await getBase64Mask()
    sdImg2ImgPrompt.mask = response
  }
  return sdImg2ImgPrompt
}

const enablePartPaint = (checked: boolean) => {
  if (!checked) {
    ctx.value.clearRect(0, 0, canvasWidth.value, canvasHeight.value)
    canvasHistory.value = []
    sdImg2ImgPrompt.mask = ''
  }
}

const addPrompt = (prompt: string, which: number) => {
  if (which == -1) {
    sdImg2ImgPrompt.negativePrompt = sdImg2ImgPrompt.negativePrompt + "," + prompt
  } else {
    sdImg2ImgPrompt.prompt = sdImg2ImgPrompt.prompt + "," + prompt
  }
}

defineExpose({
  getImg2ImgPrompt, addPrompt
})
</script>

<style>
.canvas-wrapper {
  position: relative;
}

image {
  display: block;
}

canvas {
  position: absolute;
  top: 0;
  left: 0;
}
</style>
